package cn.mengfly.algorithm.search;

import cn.mengfly.algorithm.util.ArrayUtil;

/**
 * 加权QuickUnion算法
 *
 * @author wangp
 */
public class WeightQuickUnionUf implements Uf {

    /**
     * 父连接数组，通过触点作为索引
     */
    private final int[] id;
    /**
     * 通过触点索引的各分量的大小
     */
    private final int[] sz;

    private int count;

    public WeightQuickUnionUf(int n) {
        this.count = n;
        this.id = ArrayUtil.range(n);
        this.sz = ArrayUtil.range(n);
    }

    @Override
    public void union(int p, int q) {
        int pNum = find(p);
        int qNum = find(q);
        if (pNum == qNum) {
            return;
        }
        if (sz[pNum] < sz[qNum]) {
            id[pNum] = qNum;
            sz[qNum] += sz[pNum];
        } else {
            id[qNum] = id[pNum];
            sz[pNum] += sz[qNum];
        }
        count--;

    }

    @Override
    public int find(int p) {
        while (p != id[p]) {
            p = id[p];
        }
        return p;
    }

    @Override
    public boolean connected(int p, int q) {
        return find(p) == find(q);
    }

    @Override
    public int count() {
        return count;
    }
}
